import json
import math
import random
import statistics
from time import time
import ndcg
import template
from openai import OpenAI
import asyncio


client = OpenAI(api_key="sk-d7ec8df80d9841d3a4aa46d3dbb79d7a",
                base_url="https://api.deepseek.com")


def 调用大模型(p):
    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=p,
        stream=False
    )

    return response.choices[0].message.content


async def 测试单个(d, ndcg10s=[], texts=[]):
    实际观影电影id, _ = d['target_item']

    推荐电影ids = []
    p = template.construct_prompt(d)
    text = 调用大模型(p)
    texts.append(text)
    推荐电影ids = template.parse_output(text)

    if not 推荐电影ids:
        推荐电影ids = [id for id, _ in d['candidates']]
        random.shuffle(推荐电影ids)
    ndcg10s.append(ndcg.calculate_ndcg_for_sample(推荐电影ids, 实际观影电影id, 10))


async def 测试():
    with open('val.jsonl', 'r', encoding='utf-8') as file:
        ndcg10s, texts = [], []
        tasks = []
        for line in file:
            d = json.loads(line)
            tasks.append(asyncio.create_task(测试单个(d, ndcg10s, texts)))

        await asyncio.gather(*tasks)
    return ndcg10s, texts


if __name__ == '__main__':
    start_time = time()
    ndcg10s, texts = asyncio.run(测试())
    print(f'模型得分: {statistics.mean(ndcg10s)}, {tuple(ndcg10s)}, {texts}')
    print('耗时:', time()-start_time)

    with open('out_texts/scores.txt', 'w', encoding='utf-8') as f:
        f.writelines(f'模型得分: {statistics.mean(ndcg10s)}, {tuple(ndcg10s)}')

    for text_id, text in enumerate(texts):
        with open(f'out_texts/{text_id}.md', 'w', encoding='utf-8') as f:
            f.write(text)

    # text = """
    # 数据1: [1, 2, 3],
    # 数据2: [4, 5, 6],
    # 最后的数据: [7, 8, 9]
    # """
    # print( template.parse_output(text))
    pass
